import torch
import torch.nn as nn
import torch.nn.functional as F

from models.osmnet import L2NormDense, input_norm

class UNetMatch(nn.Module):

    def __init__(self, out_ch=4) -> None:
        super().__init__()

        in_ch = 1

        self.encoder1 = nn.Sequential(
            nn.Conv2d(in_ch, 64, kernel_size=3, padding=1, bias=False), 
            nn.BatchNorm2d(64), 
            nn.ReLU(inplace=True), 
            nn.Conv2d(64, 64, kernel_size=3, padding=1,  bias=False), 
            nn.BatchNorm2d(64), 
            nn.ReLU(inplace=True), 
        )
        self.down1 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.encoder2 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, padding=1, bias=False), 
            nn.BatchNorm2d(128), 
            nn.ReLU(inplace=True), 
            nn.Conv2d(128, 128, kernel_size=3, padding=1,  bias=False), 
            nn.BatchNorm2d(128), 
            nn.ReLU(inplace=True), 
        )
        self.down2 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.encoder3 = nn.Sequential(
            nn.Conv2d(128, 256, kernel_size=3, padding=1, bias=False), 
            nn.BatchNorm2d(256), 
            nn.ReLU(inplace=True), 
            nn.Conv2d(256, 256, kernel_size=3, padding=1,  bias=False), 
            nn.BatchNorm2d(256), 
            nn.ReLU(inplace=True), 
        )

        self.up1 = nn.Sequential(
            nn.ConvTranspose2d(256, 128, kernel_size=3, 
                stride=2, padding=1, bias=False, output_padding=1), 
            nn.BatchNorm2d(128), 
            nn.ReLU(inplace=True), 
        )
        self.decoder1 = nn.Sequential(
            nn.Conv2d(256, 128, kernel_size=3, padding=1, bias=False), 
            nn.BatchNorm2d(128), 
            nn.ReLU(inplace=True), 
            nn.Conv2d(128, 128, kernel_size=3, padding=1,  bias=False), 
            nn.BatchNorm2d(128), 
            nn.ReLU(inplace=True), 
        )

        self.up2 = nn.Sequential(
            nn.ConvTranspose2d(128, 64, kernel_size=3, 
                stride=2, padding=1, bias=False, output_padding=1), 
            nn.BatchNorm2d(64), 
            nn.ReLU(inplace=True), 
        )
        self.decoder2 = nn.Sequential(
            nn.Conv2d(128, 64, kernel_size=3, padding=1, bias=False), 
            nn.BatchNorm2d(64), 
            nn.ReLU(inplace=True), 
            nn.Conv2d(64, 64, kernel_size=3, padding=1,  bias=False), 
            nn.BatchNorm2d(64), 
            nn.ReLU(inplace=True), 
        )

        self.reduce = nn.Sequential(
            nn.Conv2d(64, out_ch, kernel_size=3, padding=1, bias=False), 
            nn.BatchNorm2d(4), 
            nn.ReLU(inplace=True)
        )
        self.norm = L2NormDense()

    def forward(self, x):
        x = input_norm(x)

        ft1 = self.encoder1(x)
        ft = self.down1(ft1)
        ft2 = self.encoder2(ft)
        ft = self.down2(ft2)
        ft3 = self.encoder3(ft)

        ft = self.up1(ft3)
        ft = torch.cat([ft2, ft], dim=1)
        ft = self.decoder1(ft)
        ft = self.up2(ft)
        ft = torch.cat([ft1, ft], dim=1)
        ft = self.decoder2(ft)

        out = self.reduce(ft)
        norm_out = self.norm(out)

        return out, norm_out

if __name__ == "__main__":

    net = UNetMatch().cuda()

    img = torch.zeros((4, 1, 128, 128)).float().cuda()

    _, des = net(img)
    print(des.shape)
